import torch
import torchvision
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from model import Model



'''
使用GPU训练模型步骤：改动网络模型，损失函数，数据输入及标注即可，使用 .cuda()
1. 准备数据集：测试以及训练的数据集
2. 加载数据集
3. 创建网络模型
4. 设置损失函数
5. 准备优化器
6. 开始训练：优化器优化模型，进行测试训练
'''

# 准备数据集
train_data = torchvision.datasets.CIFAR10(root='../../../data',train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data = torchvision.datasets.CIFAR10(root='../../../data',train=False,transform=torchvision.transforms.ToTensor(),download=True)

# 创建数据加载器
train_dataLoader = DataLoader(train_data,batch_size=64)
test_dataLoader = DataLoader(test_data,batch_size=64)

# 创建网络模型
model = Model()
if torch.cuda.is_available():
    print("使用GPU")
    model = model.cuda()


# 设置损失函数
loss_fn = CrossEntropyLoss()
if torch.cuda.is_available():
    loss_fn = loss_fn.cuda()

# 准备优化器
learn_rate = 0.01
optimizer = torch.optim.SGD(model.parameters(),lr=learn_rate)

# 设置训练参数
# 训练次数和测试次数
total_train_step= 0
total_test_step= 0
# 训练轮数
epoch = 10
write = SummaryWriter('../../../logs_train_gpu')
# 训练
for i in range(epoch):
    print("-----第{}轮训练开始-----".format(i+1))
    model.train()
    for data in train_dataLoader:
        imgs,targets = data
        if torch.cuda.is_available():
            imgs = imgs.cuda()
            targets = targets.cuda()
        outputs = model(imgs)
        loss = loss_fn(outputs,targets)

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        if total_train_step % 100 == 0:
            print(("训练次数：{},loss:{}".format(total_train_step,loss.item())))
            write.add_scalar("train_loss",loss.item(),total_train_step)

    # 测试
    model.eval()
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataLoader:
            imgs,targets = data
            if torch.cuda.is_available():
                imgs = imgs.cuda()
                targets = targets.cuda()
            outputs = model(imgs)
            loss = loss_fn(outputs,targets)
            total_test_loss += loss.item()
            total_test_step += 1
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy += accuracy
    print("整体测试集上的Loss:{}".format(total_test_loss))
    test_accuracy = total_accuracy/len(test_dataLoader)
    print("整体测试集上的正确率:{}".format(test_accuracy))
    write.add_scalar("test_loss",total_test_loss,total_test_step)
    write.add_scalar("test_accuracy",test_accuracy,total_test_step)
    total_test_step += 1
    torch.save(model,"model_{}.pth".format(i))
    print("模型已保存")
write.close()


